"""
BGE-M3 API Embeddings 封装类
将本地加载的 BGE-M3 模型替换为 API 调用

参考: bge-m3 api-test.py
"""

import os
import requests
import json
import logging
from typing import List, Union
from langchain.embeddings.base import Embeddings
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

logger = logging.getLogger(__name__)


class BGEM3APIEmbeddings(Embeddings):
    """
    BGE-M3 API 嵌入模型封装类
    
    使用 ModelArts MaaS API 调用 BGE-M3 模型生成向量嵌入
    替代本地加载的 HuggingFaceEmbeddings
    """
    
    def __init__(
        self,
        api_url: str = None,
        api_key: str = None,
        model_name: str = "bge-m3",
        encoding_format: str = "float"
    ):
        """
        初始化 BGE-M3 API Embeddings
        
        Args:
            api_url: API 地址，如果未提供则从环境变量 BGE_M3_API_URL 读取
            api_key: API Key，如果未提供则从环境变量 BGE_M3_API_KEY 或 MAAS_API_KEY 读取
            model_name: 模型名称，默认为 "bge-m3"
            encoding_format: 编码格式，"float" 或 "base64"，默认为 "float"
        """
        # 从环境变量读取配置
        self.api_url = api_url or os.getenv(
            "BGE_M3_API_URL",
            "https://maas-cn-southwest-2.modelarts-maas.com/v1/infers/000e0d3c-1129-45dd-8066-976f7b755c68/v1/embeddings"
        )
        self.api_key = api_key or os.getenv("BGE_M3_API_KEY") or os.getenv("MAAS_API_KEY")
        self.model_name = model_name
        self.encoding_format = encoding_format
        
        # 验证必需参数
        if not self.api_key:
            raise ValueError(
                "API Key 未设置。请设置环境变量 BGE_M3_API_KEY 或 MAAS_API_KEY，"
                "或者在初始化时传入 api_key 参数"
            )
        
        if not self.api_url:
            raise ValueError(
                "API URL 未设置。请设置环境变量 BGE_M3_API_URL，"
                "或者在初始化时传入 api_url 参数"
            )
        
        logger.info(f"初始化 BGE-M3 API Embeddings")
        logger.info(f"API URL: {self.api_url}")
        logger.info(f"Model: {self.model_name}")
    
    def _call_api(self, texts: List[str]) -> List[List[float]]:
        """
        调用 BGE-M3 API 生成向量嵌入
        
        Args:
            texts: 文本列表
            
        Returns:
            List[List[float]]: 向量嵌入列表
        """
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        
        data = {
            "model": self.model_name,
            "input": texts,
            "encoding_format": self.encoding_format
        }
        
        try:
            response = requests.post(
                self.api_url,
                headers=headers,
                data=json.dumps(data),
                verify=False,
                timeout=60  # 60秒超时
            )
            
            response.raise_for_status()
            
            result = response.json()
            
            # 解析响应
            if isinstance(result, dict):
                if 'data' in result:
                    # OpenAI 格式响应
                    embeddings = [item['embedding'] for item in result['data']]
                elif 'embeddings' in result:
                    # 其他可能的响应格式
                    embeddings = result['embeddings']
                else:
                    raise ValueError(f"无法解析API响应: {result}")
            elif isinstance(result, list):
                # 直接返回列表格式
                embeddings = result
            else:
                raise ValueError(f"未知的API响应格式: {type(result)}")
            
            logger.debug(f"成功生成 {len(embeddings)} 个向量嵌入")
            return embeddings
            
        except requests.exceptions.RequestException as e:
            logger.error(f"API 调用失败: {str(e)}")
            raise Exception(f"BGE-M3 API 调用失败: {str(e)}")
        except Exception as e:
            logger.error(f"处理 API 响应失败: {str(e)}")
            raise
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        嵌入文档列表
        
        Args:
            texts: 文本列表
            
        Returns:
            List[List[float]]: 向量嵌入列表
        """
        if not texts:
            return []
        
        logger.debug(f"嵌入 {len(texts)} 个文档")
        
        # 批量处理，避免单次请求过大
        batch_size = 50  # 每次最多处理50个文本
        all_embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            batch_embeddings = self._call_api(batch)
            all_embeddings.extend(batch_embeddings)
        
        return all_embeddings
    
    def embed_query(self, text: str) -> List[float]:
        """
        嵌入查询文本
        
        Args:
            text: 查询文本
            
        Returns:
            List[float]: 向量嵌入
        """
        logger.debug(f"嵌入查询: {text[:50]}...")
        embeddings = self._call_api([text])
        return embeddings[0] if embeddings else []

